from flask import Flask, render_template, request, jsonify, send_file
import pandas as pd
import numpy as np
import matplotlib
matplotlib.use('Agg')  # Use non-interactive backend
import matplotlib.pyplot as plt
import matplotlib.font_manager as fm
import io
import base64
import os
import json
from decimal import Decimal, getcontext
from math import sin, cos, pi, tan, degrees
import uuid

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'data'

# 确保data文件夹存在
os.makedirs('data', exist_ok=True)

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimSun']  # 宋体
plt.rcParams['axes.unicode_minus'] = False

class MohrCircleCalculator:
    def __init__(self):
        self.circles_data = []
        self.coefficients = {}
        self.theta_result = None
        self.line_params = {}
        
    def add_circle_data(self, x, y, r, confining_pressure=None, compressive_strength=None):
        """添加圆数据"""
        circle_data = {'x': x, 'y': y, 'r': r}
        if confining_pressure is not None and compressive_strength is not None:
            circle_data['confining_pressure'] = confining_pressure
            circle_data['compressive_strength'] = compressive_strength
        self.circles_data.append(circle_data)
    
    def clear_data(self):
        """清空所有数据"""
        self.circles_data = []
        self.coefficients = {}
        self.theta_result = None
        self.line_params = {}
    
    def convert_experimental_data(self, confining_pressure, compressive_strength):
        """将实验数据转换为圆数据"""
        circles = []
        for cp, cs in zip(confining_pressure, compressive_strength):
            # 根据正确的莫尔圆理论转换
            # 圆心的坐标（在横轴上）= (抗压强度 + 围压) / 2
            # 圆的半径 = (抗压强度 - 围压) / 2
            center_x = (cs + cp) / 2
            center_y = 0
            radius = (cs - cp) / 2
            circles.append({
                'x': center_x,
                'y': center_y,
                'r': radius,
                'confining_pressure': cp,  # 保留原始围压数据
                'compressive_strength': cs  # 保留原始抗压强度数据
            })
        return circles
    
    def calculate_coefficients(self):
        """计算A、B、C、D系数"""
        if not self.circles_data:
            return None
            
        df = pd.DataFrame(self.circles_data)
        x = df['x'].to_numpy()
        y = df['y'].to_numpy()
        r = df['r'].to_numpy()
        
        x_bar = np.mean(x)
        y_bar = np.mean(y)
        r_bar = np.mean(r)
        
        A = 0.5 * np.sum(x**2 - y**2 - x_bar * x + y_bar * y)
        B = -0.5 * np.sum(2 * x * y - x_bar * y - y_bar * x)
        C = np.sum(r_bar * y - r * y)
        D = np.sum(r_bar * x - r * x)
        
        self.coefficients = {
            'A': float(A),
            'B': float(B),
            'C': float(C),
            'D': float(D),
            'x_mean': float(x_bar),
            'y_mean': float(y_bar),
            'r_mean': float(r_bar)
        }
        
        return self.coefficients
    
    def newton_iteration(self, theta_0, precision=10, max_iterations=100):
        """牛顿迭代法求解θ - 修复数据溢出问题"""
        # 设置高精度计算环境
        original_prec = getcontext().prec
        getcontext().prec = max(precision + 10, 50)  # 增加额外精度以避免舍入误差

        try:
            # 使用科学计数法处理大数值
            A = Decimal(f"{self.coefficients['A']:.15e}")
            B = Decimal(f"{self.coefficients['B']:.15e}")
            C = Decimal(f"{self.coefficients['C']:.15e}")
            D = Decimal(f"{self.coefficients['D']:.15e}")

            # 检查系数是否过大
            max_coeff = max(abs(A), abs(B), abs(C), abs(D))
            if max_coeff > Decimal('1e10'):
                # 对系数进行归一化
                scale_factor = Decimal('1e6') / max_coeff
                A *= scale_factor
                B *= scale_factor
                C *= scale_factor
                D *= scale_factor

            theta = Decimal(str(theta_0))
            iterations = []

            # 收敛阈值
            convergence_threshold = Decimal(f"1e-{precision}")

            for i in range(max_iterations):
                try:
                    # 使用更稳定的三角函数计算
                    theta_float = float(theta)

                    # 限制theta在合理范围内
                    if abs(theta_float) > 10 * pi:
                        theta_float = theta_float % (2 * pi)
                        theta = Decimal(str(theta_float))

                    sin_theta = Decimal(str(sin(theta_float)))
                    cos_theta = Decimal(str(cos(theta_float)))
                    sin_2theta = Decimal(str(sin(2 * theta_float)))
                    cos_2theta = Decimal(str(cos(2 * theta_float)))

                    # 计算函数值
                    f_theta = A * sin_2theta + B * cos_2theta + C * sin_theta + D * cos_theta

                    # 计算导数
                    f_prime = 2 * A * cos_2theta - 2 * B * sin_2theta + C * cos_theta - D * sin_theta

                    # 检查导数是否接近零
                    if abs(f_prime) < Decimal('1e-15'):
                        iterations.append({
                            'iteration': i + 1,
                            'theta': float(theta),
                            'f_theta': float(f_theta),
                            'delta': 0.0,
                            'note': '导数接近零，停止迭代'
                        })
                        break

                    # 牛顿迭代步骤
                    delta_theta = f_theta / f_prime
                    theta_new = theta - delta_theta
                    delta = abs(delta_theta)

                    iterations.append({
                        'iteration': i + 1,
                        'theta': float(theta_new),
                        'f_theta': float(f_theta),
                        'delta': float(delta)
                    })

                    # 检查收敛
                    if delta < convergence_threshold:
                        break

                    theta = theta_new

                except (OverflowError, ValueError) as e:
                    iterations.append({
                        'iteration': i + 1,
                        'theta': float(theta),
                        'f_theta': float('inf'),
                        'delta': float('inf'),
                        'note': f'数值错误: {str(e)}'
                    })
                    break

            self.theta_result = float(theta)
            return iterations

        finally:
            # 恢复原始精度
            getcontext().prec = original_prec

    def newton_single_step(self, current_theta, precision=10):
        """执行单步牛顿迭代"""
        # 设置高精度计算环境
        original_prec = getcontext().prec
        getcontext().prec = max(precision + 10, 50)

        try:
            # 使用科学计数法处理大数值
            A = Decimal(f"{self.coefficients['A']:.15e}")
            B = Decimal(f"{self.coefficients['B']:.15e}")
            C = Decimal(f"{self.coefficients['C']:.15e}")
            D = Decimal(f"{self.coefficients['D']:.15e}")

            # 检查系数是否过大并归一化
            max_coeff = max(abs(A), abs(B), abs(C), abs(D))
            if max_coeff > Decimal('1e10'):
                scale_factor = Decimal('1e6') / max_coeff
                A *= scale_factor
                B *= scale_factor
                C *= scale_factor
                D *= scale_factor

            theta = Decimal(str(current_theta))
            convergence_threshold = Decimal(f"1e-{precision}")

            try:
                theta_float = float(theta)

                # 限制theta在合理范围内
                if abs(theta_float) > 10 * pi:
                    theta_float = theta_float % (2 * pi)
                    theta = Decimal(str(theta_float))

                sin_theta = Decimal(str(sin(theta_float)))
                cos_theta = Decimal(str(cos(theta_float)))
                sin_2theta = Decimal(str(sin(2 * theta_float)))
                cos_2theta = Decimal(str(cos(2 * theta_float)))

                # 计算函数值
                f_theta = A * sin_2theta + B * cos_2theta + C * sin_theta + D * cos_theta

                # 计算导数
                f_prime = 2 * A * cos_2theta - 2 * B * sin_2theta + C * cos_theta - D * sin_theta

                # 检查导数是否接近零
                if abs(f_prime) < Decimal('1e-15'):
                    return {
                        'new_theta': float(theta),
                        'f_theta': float(f_theta),
                        'delta': 0.0,
                        'converged': True,
                        'note': '导数接近零'
                    }

                # 牛顿迭代步骤
                delta_theta = f_theta / f_prime
                theta_new = theta - delta_theta
                delta = abs(delta_theta)

                # 检查收敛
                converged = delta < convergence_threshold

                # 更新theta_result为最新值
                self.theta_result = float(theta_new)

                return {
                    'new_theta': float(theta_new),
                    'f_theta': float(f_theta),
                    'delta': float(delta),
                    'converged': converged
                }

            except (OverflowError, ValueError) as e:
                return {
                    'new_theta': float(theta),
                    'f_theta': float('inf'),
                    'delta': float('inf'),
                    'converged': False,
                    'error': str(e)
                }

        finally:
            # 恢复原始精度
            getcontext().prec = original_prec
    
    def calculate_line_parameters(self):
        """计算直线方程参数"""
        if self.theta_result is None:
            return None
            
        theta = self.theta_result
        x_mean = self.coefficients['x_mean']
        y_mean = self.coefficients['y_mean']
        r_mean = self.coefficients['r_mean']
        
        b = r_mean - x_mean * sin(theta) + y_mean * cos(theta)
        m = tan(theta)
        c = b / cos(theta)
        
        self.line_params = {
            'b': float(b),
            'slope': float(m),
            'intercept': float(c),
            'theta_degrees': float(degrees(theta)),
            'theta_radians': float(theta)
        }
        
        return self.line_params

# 全局计算器实例
calculator = MohrCircleCalculator()

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/api/upload_excel', methods=['POST'])
def upload_excel():
    """处理Excel文件上传"""
    try:
        if 'file' not in request.files:
            return jsonify({'error': '没有文件被上传'}), 400
        
        file = request.files['file']
        if file.filename == '':
            return jsonify({'error': '没有选择文件'}), 400
        
        # 保存文件
        filename = f"{uuid.uuid4()}_{file.filename}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        file.save(filepath)
        
        # 读取Excel文件
        df = pd.read_excel(filepath)
        
        # 检查列名
        if all(col in df.columns for col in ['x', 'y', 'r']):
            # 圆数据模式
            calculator.clear_data()
            for _, row in df.iterrows():
                calculator.add_circle_data(float(row['x']), float(row['y']), float(row['r']))
            data_type = 'circles'
        elif all(col in df.columns for col in ['围压', '抗压强度']):
            # 实验数据模式
            circles = calculator.convert_experimental_data(
                df['围压'].tolist(),
                df['抗压强度'].tolist()
            )
            calculator.clear_data()
            for circle in circles:
                calculator.add_circle_data(
                    circle['x'],
                    circle['y'],
                    circle['r'],
                    circle['confining_pressure'],
                    circle['compressive_strength']
                )
            data_type = 'experimental'
        else:
            return jsonify({'error': '文件格式不正确，需要包含 (x,y,r) 或 (围压,抗压强度) 列'}), 400
        
        return jsonify({
            'success': True,
            'data': calculator.circles_data,
            'data_type': data_type,
            'message': f'成功导入 {len(calculator.circles_data)} 个数据点'
        })
        
    except Exception as e:
        return jsonify({'error': f'文件处理错误: {str(e)}'}), 500

@app.route('/api/add_manual_data', methods=['POST'])
def add_manual_data():
    """添加手动输入的数据"""
    try:
        data = request.json
        data_type = data.get('type', 'circles')
        
        if data_type == 'circles':
            x = float(data['x'])
            y = float(data['y'])
            r = float(data['r'])
            calculator.add_circle_data(x, y, r)
        elif data_type == 'experimental':
            cp = float(data['confining_pressure'])
            cs = float(data['compressive_strength'])
            circles = calculator.convert_experimental_data([cp], [cs])
            calculator.add_circle_data(
                circles[0]['x'],
                circles[0]['y'],
                circles[0]['r'],
                circles[0]['confining_pressure'],
                circles[0]['compressive_strength']
            )
        
        return jsonify({
            'success': True,
            'data': calculator.circles_data,
            'data_type': data_type,
            'message': '数据添加成功'
        })
        
    except Exception as e:
        return jsonify({'error': f'数据添加失败: {str(e)}'}), 500

@app.route('/api/clear_data', methods=['POST'])
def clear_data():
    """清空所有数据"""
    calculator.clear_data()
    return jsonify({'success': True, 'message': '数据已清空'})

@app.route('/api/calculate_coefficients', methods=['POST'])
def calculate_coefficients():
    """计算A、B、C、D系数"""
    try:
        if not calculator.circles_data:
            return jsonify({'error': '没有数据可计算'}), 400

        coefficients = calculator.calculate_coefficients()
        return jsonify({
            'success': True,
            'coefficients': coefficients,
            'message': '系数计算完成'
        })

    except Exception as e:
        return jsonify({'error': f'计算系数失败: {str(e)}'}), 500

@app.route('/api/newton_iteration', methods=['POST'])
def newton_iteration():
    """牛顿迭代法求解θ"""
    try:
        data = request.json
        theta_0 = data.get('theta_0', 0)
        precision = data.get('precision', 10)

        if not calculator.coefficients:
            return jsonify({'error': '请先计算系数'}), 400

        iterations = calculator.newton_iteration(theta_0, precision)
        return jsonify({
            'success': True,
            'iterations': iterations,
            'final_theta': calculator.theta_result,
            'message': '牛顿迭代完成'
        })

    except Exception as e:
        return jsonify({'error': f'牛顿迭代失败: {str(e)}'}), 500

@app.route('/api/newton_step', methods=['POST'])
def newton_step():
    """牛顿迭代单步计算"""
    try:
        data = request.json
        current_theta = data.get('current_theta', 0)
        precision = data.get('precision', 10)

        if not calculator.coefficients:
            return jsonify({'error': '请先计算系数'}), 400

        # 执行单步牛顿迭代
        result = calculator.newton_single_step(current_theta, precision)

        return jsonify({
            'success': True,
            'new_theta': result['new_theta'],
            'f_theta': result['f_theta'],
            'delta': result['delta'],
            'converged': result['converged'],
            'message': '单步迭代完成'
        })

    except Exception as e:
        return jsonify({'error': f'单步迭代失败: {str(e)}'}), 500

@app.route('/api/calculate_line_parameters', methods=['POST'])
def calculate_line_parameters():
    """计算直线方程参数"""
    try:
        data = request.json or {}
        current_theta = data.get('current_theta')

        # 如果传入了current_theta，使用它；否则使用theta_result
        if current_theta is not None:
            calculator.theta_result = current_theta
        elif calculator.theta_result is None:
            return jsonify({'error': '请先完成牛顿迭代'}), 400

        line_params = calculator.calculate_line_parameters()
        return jsonify({
            'success': True,
            'line_params': line_params,
            'theta_used': calculator.theta_result,
            'message': f'直线参数计算完成 (θ = {calculator.theta_result:.6f})'
        })

    except Exception as e:
        return jsonify({'error': f'计算直线参数失败: {str(e)}'}), 500

@app.route('/api/export_plot/<format>')
def export_plot(format):
    """导出图片"""
    try:
        if not calculator.circles_data:
            return jsonify({'error': '没有数据可导出'}), 400

        # 获取导出选项参数
        include_envelope = request.args.get('include_envelope', 'true').lower() == 'true'
        include_legend = request.args.get('include_legend', 'true').lower() == 'true'

        # 创建matplotlib图形
        fig, ax = plt.subplots(figsize=(10, 8))

        # 绘制莫尔圆
        colors = plt.cm.tab10(np.linspace(0, 1, len(calculator.circles_data)))
        for i, circle in enumerate(calculator.circles_data):
            theta = np.linspace(0, 2*np.pi, 100)
            x = circle['x'] + circle['r'] * np.cos(theta)
            y = circle['y'] + circle['r'] * np.sin(theta)
            # 根据include_legend参数决定是否添加标签
            label = f'圆 {i+1}' if include_legend else None
            ax.plot(x, y, color=colors[i], linewidth=2, label=label)
            ax.plot(circle['x'], circle['y'], 'ko', markersize=6)

        # 添加坐标轴（不带箭头，箭头将在最后添加）
        # 添加τ=0处的横轴（更明显的样式）
        ax.axhline(y=0, color='black', linewidth=1.5, linestyle='-', alpha=1.0)

        # 添加σ=0处的纵轴（竖线）
        ax.axvline(x=0, color='black', linewidth=1.5, linestyle='-', alpha=1.0)

        # 如果有直线参数且选择包含包络线，绘制包络线
        if calculator.line_params and include_envelope:
            # 获取当前坐标轴范围
            xlim = ax.get_xlim()
            # 原始绘制范围
            original_start = xlim[0]
            original_end = xlim[1]

            # 检查最小值是否大于零，如果是则调整起点
            if original_start > 0:
                # 记录原本要绘制的起点和终点（横坐标）
                original_range = original_end - original_start
                # 将起点位置变为：0 - (终点-起点)/12
                adjusted_start = 0 - original_range / 12
                x_line = np.linspace(adjusted_start, original_end, 100)
            else:
                x_line = np.linspace(original_start, original_end, 100)

            y_line = calculator.line_params['slope'] * x_line + calculator.line_params['intercept']
            # 根据include_legend参数决定是否添加标签
            envelope_label = '包络线' if include_legend else None
            ax.plot(x_line, y_line, 'r--', linewidth=2, label=envelope_label)

            # 添加包络线公式
            theta_deg = calculator.line_params['theta_degrees']
            tau_0 = calculator.line_params['intercept']

            # 格式化公式文本，处理正负号
            if tau_0 >= 0:
                formula_text = f'τ = tan {theta_deg:.2f}° σ + {tau_0:.2f}'
            else:
                formula_text = f'τ = tan {theta_deg:.2f}° σ - {abs(tau_0):.2f}'

            # 在图的右上角添加公式（置于最顶层，避免被坐标轴箭头遮挡）
            ax.text(0.02, 0.98, formula_text,
                   transform=ax.transAxes,
                   fontsize=14,
                   fontfamily='Times New Roman',
                   bbox=dict(boxstyle='round,pad=0.8', facecolor='lightyellow', alpha=0.9, edgecolor='black'),
                   verticalalignment='top',
                   horizontalalignment='left',
                   zorder=10,
                   clip_on=False)

        ax.set_xlabel('σ (MPa)', fontsize=12, fontfamily='Times New Roman')
        ax.set_ylabel('τ (MPa)', fontsize=12, fontfamily='Times New Roman')
        ax.set_title('莫尔圆包络线图', fontsize=14, fontfamily='SimSun')
        ax.grid(True, alpha=0.3)

        # 根据include_legend参数决定是否显示图例
        if include_legend:
            ax.legend(prop={'family': 'SimSun', 'size': 10})

        ax.set_aspect('equal', adjustable='box')

        # 在绘制完所有内容后，获取最终的图形边界并添加箭头
        # 获取最终的坐标轴范围（外框线位置）
        final_xlim = ax.get_xlim()
        final_ylim = ax.get_ylim()

        # 添加横轴箭头（指向正方向），箭头位置与右外框线一致（显式降低层级，避免遮挡公式）
        ax.annotate('', xy=(final_xlim[1], 0), xytext=(final_xlim[1] - (final_xlim[1] - final_xlim[0]) * 0.05, 0),
                   arrowprops=dict(arrowstyle='->', color='black', lw=1.5),
                   zorder=2)

        # 添加纵轴箭头（指向正方向），箭头位置与上外框线一致（显式降低层级，避免遮挡公式）
        ax.annotate('', xy=(0, final_ylim[1]), xytext=(0, final_ylim[1] - (final_ylim[1] - final_ylim[0]) * 0.05),
                   arrowprops=dict(arrowstyle='->', color='black', lw=1.5),
                   zorder=2)

        # 保存图片
        img_buffer = io.BytesIO()
        if format.lower() == 'png':
            plt.savefig(img_buffer, format='png', dpi=300, bbox_inches='tight')
            mimetype = 'image/png'
        elif format.lower() == 'svg':
            plt.savefig(img_buffer, format='svg', bbox_inches='tight')
            mimetype = 'image/svg+xml'
        else:
            return jsonify({'error': '不支持的格式'}), 400

        img_buffer.seek(0)
        plt.close(fig)

        return send_file(
            img_buffer,
            mimetype=mimetype,
            as_attachment=True,
            download_name=f'mohr_circles.{format.lower()}'
        )

    except Exception as e:
        return jsonify({'error': f'导出图片失败: {str(e)}'}), 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)
